{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. Example Code to Create a composed rand cuboid group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time, sys\n",
    "import numpy as np\n",
    "sys.path.append('..')\n",
    "from pb_diff_envs.environment.static.comp_kuka_env import ComposedDynKukaEnv\n",
    "from pb_diff_envs.environment.static.comp_cuboid_group import ComposedRandCuboidGrp\n",
    "\n",
    "num_groups = 5\n",
    "num_walls_c = np.array([4, 2], dtype=np.int32)\n",
    "hr = 0.20 # 0.2\n",
    "vr_x = 0.2; vr_y = 0.2;vr_z = (0., 0.)\n",
    "se_x = 0.5; se_y = 0.5; se_z = (0.3, 0.9) # 0.4 too difficult\n",
    "hExt_range = [hr, hr, hr] # fixed-size obstacle\n",
    "seed = 333\n",
    "\n",
    "## check z\n",
    "start_end = np.array([[-se_x, se_x], [-se_y, se_y], se_z])\n",
    "half_extents_c = np.array(hExt_range)[None,].repeat(2, axis=0)\n",
    "void_range = np.array( [(-vr_x, vr_x), (-vr_y, vr_y), vr_z] ) # (0., 0.)\n",
    "\n",
    "\n",
    "cg = ComposedRandCuboidGrp('luotest-comp-kuka-vgrp-ipynb', num_groups, 2, num_walls_c, start_end, void_range, \n",
    "                            half_extents_c, seed, gen_data=True, \n",
    "                            robot_class=ComposedDynKukaEnv.robot_class,\n",
    "                            is_eval=True, debug_mode=False, GUI=False, grp_seed_diff=True, rand_iter_limit=1e5)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2. Example Code to Create a comp vgrp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import time, sys\n",
    "sys.path.append('..')\n",
    "from pb_diff_envs.environment.static.comp_kuka_env import ComposedDynKukaEnv\n",
    "\n",
    "import numpy as np\n",
    "from importlib import reload\n",
    "import matplotlib.pyplot as plt\n",
    "import pybullet as p\n",
    "import pb_diff_envs.objects.static.voxel_group as vv; reload(vv)\n",
    "from pb_diff_envs.environment.static.kuka_rand_vgrp import Kuka_VoxelRandGroupList\n",
    "no_while = True\n",
    "\n",
    "mazelist_config = dict(planner_timeout=30, \n",
    "                        interp_density=3,\n",
    "                       min_episode_distance=4*np.pi)\n",
    "\n",
    "import gym\n",
    "env_name = 'Compkuka7d-testOnly-rand-v11'\n",
    "# env_name = 'Compkuka7d-rand-nv4nv3-se0505-vr0202-hExt20-v0'\n",
    "dvg_group = gym.make(env_name, load_mazeEnv=False, gen_data=True, is_eval=True) # True\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if p.isConnected():\n",
    "    p.disconnect()\n",
    "p.isConnected()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## simple vis of env"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "if p.isConnected():\n",
    "    p.disconnect()\n",
    "env = dvg_group.create_single_env(0) # \n",
    "env.unload_env()\n",
    "env.load(GUI=False) # important not delete\n",
    "p.getContactPoints()\n",
    "p.stepSimulation()\n",
    "# env.robot.get_workspace_observation()\n",
    "plt.imshow(env.render())\n",
    "plt.show()\n",
    "# p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)  # Disable the default GUI controls\n",
    "# p.configureDebugVisualizer(p.COV_ENABLE_MOUSE_PICKING, 1)  # Enable mouse picking for camera control\n",
    "p.addUserDebugLine([0, 0, 0], [1, 0, 0], lineColorRGB=[1, 0, 0], lineWidth=6)\n",
    "p.addUserDebugLine([0, 0, 0], [0, 1, 0], lineColorRGB=[0, 1, 0], lineWidth=6)\n",
    "p.addUserDebugLine([0, 0, 0], [0, 0, 1], lineColorRGB=[0, 0, 1], lineWidth=6)\n",
    "while True:\n",
    "    p.stepSimulation()\n",
    "    # p.performCollisionDetection()\n",
    "    c_pts = p.getContactPoints()\n",
    "    # print(p.get)\n",
    "    time.sleep(0.1)\n",
    "    if c_pts is not None and len(c_pts) > 0:\n",
    "        # print(p.getContactPoints())\n",
    "        pass\n",
    "        # time.sleep(2)\n",
    "        # p.resetJointStatesMultiDof(env.robot_id, np.arange(7), [[0.]]*7)\n",
    "    if no_while:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample a problem/episode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from utils.utils import DotDict\n",
    "from tqdm import tqdm\n",
    "env.load(GUI=False) # False\n",
    "prev_pose = np.array([0.0] * 7)\n",
    "env.robot.set_config(prev_pose)\n",
    "solutions = []\n",
    "# n_traj = 25\n",
    "n_traj = 5 # 20\n",
    "for i in tqdm(range(n_traj)):\n",
    "    solution, _ = env.sample_1_episode(prev_pose)\n",
    "    solution = np.array(solution)\n",
    "    prev_pose = solution[-1]\n",
    "    solutions.append(solution)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import pybullet as p\n",
    "import pybullet_data\n",
    "import numpy as np\n",
    "from utils.utils import save_gif\n",
    "from IPython.display import HTML\n",
    "import base64\n",
    "from objects.trajectory import WaypointLinearTrajectory, WaypointProportionTrajectory\n",
    "from utils.kuka_utils_luo import SolutionInterp\n",
    "# Visualization\n",
    "from importlib import reload\n",
    "import utils.robogroup_utils_luo  as rul; reload(rul)\n",
    "from utils.robogroup_utils_luo import robogroup_visualize_traj_luo\n",
    "from utils.kuka_utils_luo import visualize_kuka_traj_luo\n",
    "\n",
    "sol_interp = SolutionInterp(density=mazelist_config['interp_density'])\n",
    "sol_i = solutions[0] # which motion trajectory to visualize\n",
    "result = DotDict(solution=sol_i)\n",
    "traj_vis = sol_interp(result.solution,)\n",
    "\n",
    "\n",
    "gifs, ds, vis_dict = visualize_kuka_traj_luo(env, traj_vis, is_debug=False)\n",
    "\n",
    "\n",
    "gif_name = f'./test_comp_kuka.gif'\n",
    "save_gif(gifs, gif_name, duration=ds)\n",
    "b64 = base64.b64encode(open(gif_name, 'rb').read()).decode('ascii')\n",
    "display(HTML(f'<img src=\"data:image/gif;base64,{b64}\" />'))\n",
    "\n",
    "print(traj_vis.shape)\n",
    "print(result.solution.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# while True:\n",
    "#     time.sleep(0.1)\n",
    "#     p.stepSimulation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "vscode": {
   "interpreter": {
    "hash": "c1daed7e4d34e14fbb9bb260207f2e4074aea67593e999acdc484a19561ddc79"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
